# 处理数据集，将数据集分类保存
import pickle


# 定义一个辅助函数，用于读取CIFAR-10的数据文件
# 它接受一个文件名作为参数，并返回解析后的字典
def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


# 包含一个CIFAR-10数据集中标签
label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]

import glob
import numpy as np
import cv2
import os

# 通过glob.glob函数获取所有以test_batch开头的文件路径列表
train_list = glob.glob("./CIFAR10/test_batch*")
print(train_list)
save_path = "./test"

for l in train_list:
    print(l)
    # 获取一个字典对象
    l_dict = unpickle(l)
    # print(l_dict)
    # 分别获得图像数据、标签、文件名
    print(l_dict.keys())

    for im_idx, im_data in enumerate(l_dict[b'data']):
        '''
            im_label:保存图像的标签
            im_name: 保存图像的文件名
            im_data: 保存图像的像素数据，它是一个长度为3072的一维数组。
            im_label_name: 根据标签值从label_name列表中获取对应的类别名称
        '''
        # b'**'代表的是字节字符串，以字节为单位的字符串
        im_label = l_dict[b'labels'][im_idx]
        im_name = l_dict[b'filenames'][im_idx]
        print(im_label, im_name, im_data)
        im_label_name = label_name[im_label]
        #  im_data被重塑为一个3x32x32的数组，表示图像的三个颜色通道和32x32的尺寸。
        im_data = np.reshape(im_data, [3, 32, 32])
        # 使用np.transpose函数重新排列数组的维度，使得颜色通道成为最后一个维度。
        im_data = np.transpose(im_data, (1, 2, 0))

        # cv2.imshow("im_data",cv2.resize(im_data, (200, 200)))
        # cv2.waitKey(0)

        if not os.path.exists("{}/{}".format(save_path,
                                             im_label_name)):
            os.mkdir("{}/{}".format(save_path, im_label_name))

        cv2.imwrite("{}/{}/{}".format(save_path,
                                      im_label_name,
                                      im_name.decode("utf-8")), im_data)
